package algorithms;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Random;

import framework.Simplex;
import framework.State;

public class PolicyIteration {
	private static boolean unchanged = false;
	private static double gamma = 0.9, errorP = 0.1;
	private static Map<String, String> policy = new HashMap<String, String>();
	private static Map<String, Double> utilities = new HashMap<String, Double>();

	/**
	 * 
	 */
	public static void init() {
		// Setup random policy
		Random rnd = new Random();
		for (String s : State.allStates.keySet()) {
			List<String> transitions = State.getPossibleTransitions(State.allStates.get(s));
			if (transitions.size() != 0) {
				policy.put(s, transitions.get(rnd.nextInt(transitions.size())));
			}
		}

		policy.put(State.toString(new double[] { 2.0, 2.0, 1.0, 2.0, 0.0 }),
				State.toString(new double[] { 2.0, 2.0, 1.0, 2.0, 0.0 }));
	}

	/**
	 * 
	 */
	public static void iterate() {
		while (!unchanged) {
			unchanged = true;
			utilities.clear();
			// Calculate utility values
			String[] states = State.allStates.keySet().toArray(new String[] {});
			// A
			double[][] A = new double[State.allStates.size()][State.allStates.size()];
			//
			for (int i = 0; i < states.length; i++) {
				for (int j = 0; j < states.length; j++) {
					String newState = policy.get(states[i]);
					
					int move = 0;
					if(newState.charAt(1) != states[i].charAt(1)) {
						move = 1;
					}
					
					double prob;
					double[] a = State.allStates.get(states[i]);
					List<String> transitions = State.getPossibleTransitions(a);
					int noTransitions = transitions.size();
					if (states[j].equalsIgnoreCase(newState)) {
						prob = 0.9;
					} else if (transitions.contains(states[j]) && states[j].charAt(move) != states[i].charAt(move)) {
						// Prevent infinity
						if (noTransitions != 0) {
							prob = 0.1;
						} else {
							prob = 0;
						}
					} else {
						prob = 0;
					}
					if (i == j) {
						// Op diagonaal is prob altijd 0
						A[i][j] = 1;
					} else {
						A[i][j] = -gamma * prob;
					}
				}
			}
			// B
			double[] B = new double[states.length];
			for (int i = 0; i < states.length; i++) {
				B[i] = State.getValue(State.allStates.get(states[i]));
			}
			// C
			double[] C = new double[states.length];
			for (int i = 0; i < states.length; i++) {
				C[i] = 1;
			}
			Simplex simplex = new Simplex(A, B, C);
			double[] x = simplex.primal();
			for (int i = 0; i < x.length; i++) {
				utilities.put(states[i], x[i]);
				State.allStates.get(states[i])[4] = x[i];
			}
			//
			for (String s : State.allStates.keySet()) {
				double[] s2 = getMaxActionByUtility(State.allStates.get(s));
				if (!policy.get(s).equalsIgnoreCase(State.toString(s2))) {
					policy.put(s, State.toString(s2));
					unchanged = false;
				}
			}
		}
	}

	/**
	 * 
	 * @param s
	 * @return
	 */
	private static double[] getMaxActionByUtility(double[] s) {
		double maxVal = -1000., current = .0;
		double[] maxAction = s;
		int move;

		List<String> states = State.getPossibleTransitions(s);
		if (states.size() == 0)
			// Absorbing state
			return s;
		//
		// System.out.println(State.toVisualString(s));
		for (String state : states) {
			double[] newState = State.allStates.get(state);
			//System.out.println(State.toVisualString(newState));
			if (newState[0] != s[0]) {
				move = 0; // A Moved
			} else {
				move = 1; // B Moved
			}
			current = State.getValue(newState) + (gamma * sumStates(states, newState, move));
			if (current > maxVal) {
				maxVal = current;
				maxAction = newState;
			}
		}
		return maxAction;
	}

	/**
	 * 
	 * @param states
	 * @param currentState
	 * @return
	 */
	private static double sumStates(List<String> states, double[] currentState, int move) {
		double sum = 0., prob = 0;
		for (String state : states) {
			double[] newState = State.allStates.get(state);
			if (newState[move] == currentState[move]) {
				continue; // Not the move we are looking for
			}
			if (currentState.equals(newState)) {
				prob = 1. - errorP;
			} else {
				prob = errorP;
			}

			sum += prob * newState[4];
		}
		return sum;
	}

	/**
	 * 
	 * @param s
	 * @return
	 */
	private static double calculateUtility(double[] s) {
		int move;
		// If we already calculated it
		if (utilities.get(State.toString(s)) != null) {
			return utilities.get(State.toString(s));
		}
		//
		double value = .0;
		// Follow up state according to policy
		double[] newState = State.allStates.get(policy.get(State.toString(s)));
		// Possible transitions from this state, regardless of policy
		List<String> states = State.getPossibleTransitions(s);
		if (states.size() != 0 && newState != null) {
			if (newState[0] != s[0]) {
				move = 0; // A Moved
			} else {
				move = 1; // B Moved
			}
			value = State.getValue(newState) + gamma * sumStates(states, newState, move);
		} else {
			value = State.getValue(s);
		}
		// Remember it
		utilities.put(State.toString(s), value);
		return value;
	}

	/**
	 * 
	 * @return
	 */
	public static String printPolicy() {
		StringBuilder builder = new StringBuilder();

		builder.append("----Policy-----\n");
		for (double[] s : State.allStates.values()) {
			//builder.append(State.toString(s));
			builder.append(State.toVisualString(s, true));
			builder.append("-->\n");
			//builder.append(policy.get(State.toString(s)));
			builder.append(State.toVisualString(State.allStates.get(policy.get(State.toString(s))), true));
			builder.append("########\n");
		}
		return builder.toString();
	}
}
